# from backtrader import indicator
import pandas as pd
from pymongo import MongoClient
from datetime import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import numpy as np


# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    '''#平滑异同移动平均线MACD
            DIF(蓝线): 计算12天平均和26天平均的差，公式：EMA(C,12)-EMA(c,26)
        Signal(DEM或DEA或MACD) (红线): 计算macd9天均值，公式：Signal(DEM或DEA或MACD)：EMA(MACD,9)
            Histogram (柱): 计算macd与signal的差值，公式：Histogram：MACD-Signal
            period_me1=12
            period_me2=26
            period_signal=9
            macd = ema(data, me1_period) - ema(data, me2_period)
            signal = ema(macd, signal_period)
            histo = macd - signal
        '''


    #初始化部分指标
    def __init__(self):
        #调用用户初始化代码
        self.dataclose = self.datas[0].close
    
        self.me1 =  bt.indicators.ExponentialMovingAverage(self.datas[0].close,period=12)
        self.me2 =  bt.indicators.ExponentialMovingAverage(self.datas[0].close,period=26)

        self.macd = self.me1 - self.me2
        self.signal =  bt.indicators.ExponentialMovingAverage(self.macd,period=9)
        
        #指标进行注册
        self.indicators = {}
        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None
        
    # 策略核心代码
    def next(self):

        #调用用户策略
        if self.order:
            return
        self.cash = self.get_cash()   #获取现金
        self.dates = self.datas[0].datetime.date(0).isoformat()  #当前日期

        
        #没有持仓
        if not self.position:
            # pass
            condition1 = self.macd[-1] - self.signal[-1]
            condition2 = self.macd[0] - self.signal[0]

            if condition1 < 0 and condition2 > 0:
                self.buy()
        # 有持仓
        else:
            # pass
            condition = (self.dataclose[0] - self.dataclose[-1]) / self.dataclose[0]

            if condition > 0.1 or condition < -0.1:
                self.sell()

        
    
    #买入
    def open_buy(self,size=None):
        self.order = self.buy(size=size)
    #卖出
    def open_sell(self,size=None):
        self.order = self.sell(size=size)
    #平仓
    def open_close(self):
        self.order = self.close()
    #取消订单
    def open_cancel(self,size=None):
        self.order = self.cancel(size=size)
    #获取资金
    def get_cash(self):
        return self.broker.getvalue()
    #获取某个时间段内的数据
    # data  获取指定列的数据
    # ago 0：从当前时间点往前取数据（默认为0）
        # 1：从当前时间点的前一天开始
    #size  时间点数据（默认为10）
    def optainData(self,data,ago=0,size = 10):
        return np.array(data.get(ago=ago, size=size))



    # 日志输出
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')


def main(execution_id, adjustflag, code, start, end='', startcash=1000000, com=0.0015, qts=100):

    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.t_strategy_execution_result
    mydict = {'_id': execution_id}
    s = mycol.find(mydict).count()
    if s > 0:
        return execution_id

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    df = selectSotck(adjustflag, code, start, end)
    df.index = pd.to_datetime(df.date)
    df = df[['open', 'high', 'low', 'close', 'volume']]
    # 将数据加载至回测系统
    data = bt.feeds.PandasData(dataname=df)

    cerebro.adddata(data)

    print(data)
    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # cerebro.broker.setcommission(PerOrder(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    # cerebro.broker.set_coc(True)
    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run(maxcpus=1)
    print('期末总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.plot()


# 时间处理
def loadDate(str):
    date = datetime.strptime(str, '%Y%m%d').strftime('%Y-%m-%d')
    return (date)


# 数据查询
def selectSotck(adjustflag,code, startDate, endDate):
    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.testKline

    myquery = {
        "adjustflag": adjustflag,  # 复权类型
        "code": code,  # 股票代码
        "date": {'$gt': startDate, '$lt': endDate}  # 时间区间
    }

    mydoc = mycol.find(myquery)
    data_list = []
    for x in mydoc:
        data_list.append(x)
    fields = ['date', 'open', 'close', 'high', 'low', 'volume', 'code']
    df = pd.DataFrame(data_list, columns=fields)
    df['date'] = df['date'].map(loadDate)
    test_list = ['open', 'close', 'high', 'low', 'volume']
    df[test_list] = df[test_list].astype('float64')
    return df

run = main(
    execution_id='1',
    adjustflag = '3',  #复权类型 3：不复权 1：后复权；2：前复权
    code='sh.603816',  #股票池
    start='20190101',    #开始时间
    end = '20210101',   #结束时间
    startcash = 10000,  #资金
    com = 0.0003   #手续费
)